import os
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import matplotlib.pyplot as plt
import random
import torch.nn.functional as F
import math


class RibFracSegmentationDataset(Dataset):
    
    def __init__(self, data_dir, transform=None, mode='train'):

        self.data_dir = data_dir
        self.transform = transform
        self.mode = mode
        
        metadata_path = os.path.join(data_dir, 'metadata', 'metadata.csv') #此处有一个问题，文件加载后的metadata有patch命名规则
        # 如果没有metadata.csv文件
        if os.path.exists(metadata_path):
            self.metadata = pd.read_csv(metadata_path)
        else:
            image_dir = os.path.join(data_dir, 'images')
            self.image_files = sorted([f for f in os.listdir(image_dir) if f.endswith('.npy')])
            self.metadata = pd.DataFrame({
                'ct_path': [os.path.join('images', f) for f in self.image_files],
                'mask_path': [os.path.join('labels', f.replace('_ct.npy', '_mask.npy')) 
                                for f in self.image_files]
            })
        
        print(f"加载{mode}数据集: {len(self.metadata)}个样本")
    
    def __len__(self):
        return len(self.metadata)
    
    def __getitem__(self, idx):

        # print(f"Sample {idx} - Random flip: {random.random() < 0.5}")
        row = self.metadata.iloc[idx]
        ct_id = str(row['public_id']) + '_' + str(row['label_id'])

        ct_path = os.path.join(self.data_dir, row['ct_path'])
        mask_path = os.path.join(self.data_dir, row['mask_path'])

        ct_block = np.load(ct_path)  # (D, H, W)
        mask_block = np.load(mask_path)  # (D, H, W)

        min_val = np.min(ct_block)
        max_val = np.max(ct_block)
        
        if max_val - min_val <= 1e-6:
            if max_val != 0:
                ct_block = ct_block / max_val  # 避免除以0
            else:
                ct_block = np.zeros_like(ct_block)
        else:
            ct_block = (ct_block - min_val) / (max_val - min_val)
        
        ct_block = ct_block[np.newaxis, ...]
        mask_block = mask_block[np.newaxis, ...]
        
        ct_tensor = torch.from_numpy(ct_block).float()
        mask_tensor = torch.from_numpy(mask_block).float()
        if self.transform:
            # CT和掩码堆叠在一起进行变换
            combined = torch.cat([ct_tensor, mask_tensor], dim=0)
            combined = self.transform(combined)
            
            ct_tensor = combined[0:1]
            mask_tensor = combined[1:2]

        mask_tensor = (mask_tensor > 0.5).float()
        
        if self.mode == 'test':
            return ct_tensor, mask_tensor, ct_id
        else:
            return ct_tensor, mask_tensor

class Random3DFlip:
    """3D随机翻转"""
    def __init__(self, p=0.5):
        self.p = p
        
    def __call__(self, x):
        if random.random() < self.p:
            axis = random.choice([1, 2, 3])  # 1:D, 2:H, 3:W
            return torch.flip(x, dims=[axis])
        return x

class Random3DRotation:
    """3D随机旋转 (90度倍数)"""
    def __init__(self, p=0.5):
        self.p = p
        
    def __call__(self, x):
        if random.random() < self.p:
            k = random.randint(0, 3)               # 旋转次数
            plane = random.choice([(1, 2), (1, 3), (2, 3)])  # 旋转平面
            return torch.rot90(x, k, dims=plane)
        return x

def get_data_transforms(mode):
    if mode == 'train':
        return transforms.Compose([
            Random3DFlip(p=0.3),
            Random3DRotation(p=0.3),
        ])
    else:
        return None

def create_data_loaders(train_dir, val_dir, test_dir, batch_size=8, num_workers=4):
    train_dataset = RibFracSegmentationDataset(
        train_dir, 
        transform=get_data_transforms('train'),
        mode='train'
    )
    
    val_dataset = RibFracSegmentationDataset(
        val_dir, 
        transform=get_data_transforms('val'),
        mode='val'
    )
    
    test_dataset = RibFracSegmentationDataset(
        test_dir, 
        transform=get_data_transforms('test'),
        mode='test'
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )
    
    return train_loader, val_loader, test_loader

def create_test_data_loaders(test_dir, batch_size=8, num_workers=4):
    test_dataset = RibFracSegmentationDataset(
        test_dir, 
        transform=get_data_transforms('test'),
        mode='test'
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )    
    return test_loader
def visualize_batch(images, masks, num_samples=4):
    images = images.detach().cpu().numpy()
    masks = masks.detach().cpu().numpy()
    
    batch_size = images.shape[0]
    num_samples = min(num_samples, batch_size)
    indices = np.random.choice(batch_size, num_samples, replace=False)
    
    fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5*num_samples))
    
    for i, idx in enumerate(indices):
        # 获取中间切片，只查看中间切片层
        depth = images.shape[2]
        slice_idx = depth // 2
        ct_slice = images[idx, 0, slice_idx]
        mask_slice = masks[idx, 0, slice_idx]
        overlay = np.zeros((*ct_slice.shape, 3))
        overlay[..., 0] = ct_slice  # 红色通道: CT
        overlay[..., 1] = mask_slice  # 绿色通道: 骨折区域
        
        if num_samples == 1:
            ax_ct, ax_mask, ax_overlay = axes
        else:
            ax_ct, ax_mask, ax_overlay = axes[i]
        
        ax_ct.imshow(ct_slice, cmap='gray')
        ax_ct.set_title(f'sample {idx+1}: CT slice')
        ax_ct.axis('off')
        
        ax_mask.imshow(mask_slice, cmap='viridis')
        ax_mask.set_title(f'frac mask')
        ax_mask.axis('off')
        
        ax_overlay.imshow(overlay)
        ax_overlay.set_title(f'mask view')
        ax_overlay.axis('off')
    
    plt.tight_layout()
    plt.savefig('train_unet/dataloader.png')
    plt.show()


if __name__ == "__main__":
    dataset_root = './ribfrac-dataset'
    train_dir = os.path.join(dataset_root, 'train')
    val_dir = os.path.join(dataset_root, 'val')
    test_dir = os.path.join(dataset_root, 'test')
    
    # 创建数据加载器
    train_loader, val_loader, test_loader = create_data_loaders(
        train_dir, val_dir, test_dir, batch_size=4
    )
    
    for images, masks in train_loader:
        print(images.shape, masks.shape)
        visualize_batch(images, masks)
        break
    